from rest_framework.throttling import SimpleRateThrottle
from django.core.cache import cache as default_cache
from api.extension import return_code
from rest_framework import exceptions
from rest_framework import status


class ThrottledException(exceptions.APIException):
    status_code = status.HTTP_429_TOO_MANY_REQUESTS
    default_code = 'throttled'


class CloudCreateThrottle(SimpleRateThrottle):
    cache = default_cache
    scope = "user"
    cache_format = "throttle_%(scope)s_%(ident)s"
    THROTTLE_RATES = {"user": "4/1m"}

    def parse_rate(self, rate):
        if rate is None:
            return None, None
        num, period = rate.split('/')
        num_request = int(num)
        duration = {"s": 1, 'm': 60, "h": 3600, "d": 86400}[period[-1]]
        count = int(period[0:-1])
        return num_request, duration * count

    def get_cache_key(self, request, view):
        if request.data.get("username"):
            ident = request.data.get("username")
        if request.data.get("phone"):
            ident = request.data.get("phone")
        if request.user.pk:
            ident = request.user.pk
        return self.cache_format % {'scope': self.scope, 'ident': ident}

    def throttle_failure(self):
        wait = self.wait()
        detail = {
            "code": return_code.TOO_MANY_REQUESTS,
            "data": "访问频率限制",
            "detail": "需等待{}秒后才能访问".format(int(wait))
        }
        raise ThrottledException(detail)

    def throttle_success(self):
        return True

    def done(self):
        self.history.insert(0, self.now)
        self.cache.set(self.key, self.history, self.duration)
